package cn.texous.demo.dnw.utils;

import cn.texous.demo.dnw.common.constant.RedisKeyEnum;
import cn.texous.demo.dnw.common.constant.TokenEnum;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.data.redis.core.RedisCallback;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.Collections;

/***/
@Component
public class RateLimitClient {

    @Autowired
    StringRedisTemplate stringRedisTemplate;
    @Qualifier("ratelimitLua")
    @Resource
    RedisScript<Long> ratelimitLua;

    @Qualifier("ratelimitInitLua")
    @Resource
    RedisScript<Long> ratelimitInitLua;

    @Qualifier("ratelimitByIdLua")
    @Resource
    RedisScript<Long> ratelimitByIdLua;

    /***/
    public TokenEnum initToken(String key) {
        TokenEnum token = TokenEnum.SUCCESS;
        Long currMillSecond = stringRedisTemplate.execute(
                (RedisCallback<Long>) redisConnection -> redisConnection.time()
        );

        /*
         * last_mill_second 最后时间毫秒
         * curr_permits 当前可用的令牌
         * max_burst 令牌桶最大值
         * rate 每秒生成几个令牌
         * app 应用
         *
         * redis.pcall("HMSET",KEYS[1],
         "last_mill_second",ARGV[1],
         "curr_permits",ARGV[2],
         "max_burst",ARGV[3],
         "rate",ARGV[4],
         "app",ARGV[5])
         */
        Long accquire = stringRedisTemplate.execute(ratelimitInitLua,
                Collections.singletonList(getKey(key)),
                currMillSecond.toString(), "100", "100000", "10", "skynet");
        if (accquire == 1) {
            token = TokenEnum.SUCCESS;
        } else if (accquire == 0) {
            token = TokenEnum.SUCCESS;
        } else {
            token = TokenEnum.FAILED;
        }
        return token;
    }

    public Long accquireTokenCounter(String ip) {
        return stringRedisTemplate.execute(ratelimitByIdLua,
                Collections.singletonList(ip), "10", "1");
    }

    /**
     * 获得key操作
     *
     * @param key key
     * @return
     */
    public TokenEnum accquireToken(String key) {
        return accquireToken(key, 1);
    }

    /***/
    public TokenEnum accquireToken(String key, Integer permits) {
        TokenEnum token = TokenEnum.SUCCESS;
        Long currMillSecond = stringRedisTemplate.execute(
                (RedisCallback<Long>) redisConnection -> redisConnection.time()
        );

        Long accquire = stringRedisTemplate.execute(ratelimitLua,
                Collections.singletonList(getKey(key)),
                permits.toString(), currMillSecond.toString());
        if (accquire == 1) {
            token = TokenEnum.SUCCESS;
        } else {
            token = TokenEnum.FAILED;
        }
        return token;
    }

    public boolean exists(String key) {
        return stringRedisTemplate.hasKey(key);
    }

    public boolean hexists(String key, String field) {
        return stringRedisTemplate.opsForHash().hasKey(key, field);
    }

    public String getKey(String key) {
        return RedisKeyEnum.RATELIMIT_KEY_PREFIX.getKey() + key;
    }

}
